import xmlrpc.client
from concurrent.futures import ThreadPoolExecutor
import typing as tp
import itertools

class WikidataQueryClient:
    def __init__(self, url: str):
        self.url = url
        self.server = xmlrpc.client.ServerProxy(url)

    def label2qid(self, label: str) -> tp.List[str]:
        return self.server.label2qid(label)

    def label2pid(self, label: str) -> tp.List[str]:
        return self.server.label2pid(label)

    def pid2label(self, pid: str) -> tp.List[str]:
        return self.server.pid2label(pid)

    def qid2label(self, qid: str) -> tp.List[str]:
        return self.server.qid2label(qid)

    def get_all_relations_of_an_entity(self, qid: str) -> tp.Dict:
        return self.server.get_all_relations_of_an_entity(qid)

    def get_tail_entities_given_head_and_relation(self, qid: str, pid: str) -> tp.Dict:
        return self.server.get_tail_entities_given_head_and_relation(qid, pid)

    def get_tail_values_given_head_and_relation(self, qid: str, pid: str) -> tp.List[str]:
        return self.server.get_tail_values_given_head_and_relation(qid, pid)

    def get_external_id_given_head_and_relation(self, qid: str, pid: str) -> tp.List[str]:
        return self.server.get_external_id_given_head_and_relation(qid, pid)

    def mid2qid(self, mid: str) -> tp.List[str]:
        return self.server.mid2qid(mid)


class MultiServerWikidataQueryClient:
    def __init__(self, server_addrs: tp.List[str]):
        self.clients = [WikidataQueryClient(addr) for addr in server_addrs]
        self.executor = ThreadPoolExecutor(max_workers=len(self.clients))

    def query_all(self, method: str, *args):
        futures = [
            self.executor.submit(getattr(client, method), *args)
            for client in self.clients
        ]
        is_dict_return = method in [
            "get_all_relations_of_an_entity",
            "get_tail_entities_given_head_and_relation",
        ]

        results = [f.result() for f in futures]

        real_results = set() if not is_dict_return else {"head": [], "tail": []}
        for res in results:
            if res == "Not Found!":
                continue
            elif isinstance(res, list):
                real_results.update(res)
            elif is_dict_return:
                real_results["head"].extend(res.get("head", []))
                real_results["tail"].extend(res.get("tail", []))
            else:
                real_results.add(res)

        return real_results if real_results else "Not Found!"

    # 封装常用方法
    def label2qid(self, label: str):
        return self.query_all("label2qid", label)

    def label2pid(self, label: str):
        return self.query_all("label2pid", label)

    def pid2label(self, pid: str):
        return self.query_all("pid2label", pid)

    def qid2label(self, qid: str):
        return self.query_all("qid2label", qid)

    def get_tail_values_given_head_and_relation(self, qid: str, pid: str):
        return self.query_all("get_tail_values_given_head_and_relation", qid, pid)

    def get_all_relations_of_an_entity(self, qid: str):
        return self.query_all("get_all_relations_of_an_entity", qid)

    def mid2qid(self, mid: str):
        return self.query_all("mid2qid", mid)

# Example usage
if __name__ == "__main__":
    # Step 1: 加载 server 地址列表
    with open("server_urls.txt", "r") as f:
        server_addrs = [line.strip() for line in f.readlines()]

    # Step 2: 创建客户端
    client = MultiServerWikidataQueryClient(server_addrs)

    # Step 3: 访问接口
    print("label2qid('Google') =>", client.label2qid("Google"))
    print("qid2label('Q95') =>", client.qid2label("Q95"))
    print("get_tail_values_given_head_and_relation('Q18615057', 'P249') =>",
          client.get_tail_values_given_head_and_relation("Q18615057", "P108"))
    print(client.label2pid("employer"))
    # print(client.get_all_relations_of_an_entity("Q95"))

